import sys
from posixpath import join
from os import path

import cupy as cp

import muse_main
import vecmap_main_torch

class Task:
    def __init__(self):
        return

if __name__ == '__main__':
    src = 'vi'
    trg = 'en'
    exp_name = 'Oct18_rawConcat'
    exp_id = src + '-' + trg
    exp_path = path.join('dumped', exp_name, exp_id)
    vecmap_params = [
            "--unsupervised",
            # "dumped/Seg04_lowcutoff/vi-en/vectors-vi.txt", "dumped/Sep06_stepcutoff/vi-en/vectors-en.txt",
            "data/wiki." + src + ".vec", "data/wiki." + trg + ".vec", 
            path.join(exp_path, src + ".mapped_emb.txt"), path.join(exp_path, trg + ".mapped_emb.txt"),
            "--verbose",
            "--direction", "forward",
            "--validation", "data/crosslingual/dictionaries/" + src + '-' + trg + ".5000-6500.txt",
            "--batch_size", "4000",
            "--vocabulary_cutoff", "20000",
            "--log", path.join(exp_path, 'log.txt'),
            "--seed", "0",
            "--cuda",
    ]
    muse_sdi_params = [
            "--src_lang", src,
            "--tgt_lang", trg,
            "--src_emb", "data/wiki." + src + ".vec",
            "--tgt_emb", "data/wiki." + trg + ".vec",
            "--n_refinement", "0",
            "--max_vocab", "20000",
            "--dis_most_frequent", "7500",
            "--n_adversarials", "1",
            # "--n_epochs", "1",
            # "--epoch_size", "10000",
            # "--map_optimizer", "sgd,lr=0.02",
            # "--dis_optimizer", "sgd,lr=0.02",
            # "--batch_size", "128",
            # "--exp_id", "jul29_renorm",
            "--exp_name", exp_name, 
            "--exp_id", exp_id, 
            "--normalize_embeddings", "renorm,center,renorm",
            "--seed", "0",
            # "--load_temp_model", "~/Documents/MI_Tlab/CUWTr/dumped/debug/jul25_19_32/"
            # "--dico_eval", "./datasets/crosslingual/dictionaries/en-es.txt"
            "--cuda",
            "--export", "",
            "--return_GAN_xwzw",
    ]
    t = Task()
    sys.argv.extend(muse_sdi_params)
    t.gan_initialized_space = muse_main.main()
    sys.argv = sys.argv[:-len(muse_sdi_params)]
    sys.argv.extend(vecmap_params)
    vecmap_main_torch.main(t)
